import asyncio
from typing import List

from llama_index.core.agent.workflow import  FunctionAgent
from llama_index.core.base.llms.types import ChatMessage
from llama_index.core.storage.chat_store.sql import SQLAlchemyChatStore
from llama_index.core.tools import QueryEngineTool
from llama_index.core.vector_stores import SimpleVectorStore
from llama_index.core.schema import  TextNode
from llama_index.core import Settings, SimpleKeywordTableIndex, SummaryIndex, get_response_synthesizer, StorageContext, \
    VectorStoreIndex
from llama_index.embeddings.zhipuai import ZhipuAIEmbedding
from llama_index.core.graph_stores import SimplePropertyGraphStore
from llama_index.core.schema import Document
from pydantic import BaseModel

embed_model = ZhipuAIEmbedding(
    model="embedding-2",
    api_key="f387f5e4837d4e4bba6d267682a957c9.PmPiTw8qVlsI2Oi5"
    # With the `embedding-3` class
    # of models, you can specify the size
    # of the embeddings you want returned.
    # dimensions=1024
)
Settings.embed_model=embed_model

from llama_index.llms.deepseek import DeepSeek

llm = DeepSeek(model="deepseek-chat", api_key="sk-605e60a1301040759a821b6b677556fb")
Settings.llm = llm



from llama_index.core.node_parser import HTMLNodeParser
from llama_index.readers.file import FlatReader
from pathlib import Path
'''
# 加载HTML文档
html_docs = FlatReader().load_data(Path("1.html"))

# 创建解析器实例（指定解析<p>和<h1>标签）
parser = HTMLNodeParser(tags=["nav", "main","footer"])

# 从文档获取节点列表
nodes = parser.get_nodes_from_documents(html_docs)

# 输出节点内容
for node in nodes:
    print(f"Node ID: {node.node_id}")
    print(f"Content: {node.text[:50]}...")
    print("-" * 30)

from llama_index.core.node_parser.file.html import HTMLNodeParser
from llama_index.core.node_parser.file.json import JSONNodeParser
from llama_index.core.node_parser.file.markdown import MarkdownNodeParser
from llama_index.core.node_parser.file.simple_file import SimpleFileNodeParser


from llama_index.core import SimpleDirectoryReader
from llama_index.core.node_parser import MarkdownNodeParser

reader = SimpleDirectoryReader("./docs")
docs=reader.load_data()
print(docs)
parser = MarkdownNodeParser()

nodes = parser.get_nodes_from_documents(reader.load_data())

print(nodes)
'''



from llama_index.core.node_parser.file.html import HTMLNodeParser
from llama_index.core.node_parser.file.json import JSONNodeParser
from llama_index.core.node_parser.file.markdown import MarkdownNodeParser
from llama_index.core.node_parser.file.simple_file import SimpleFileNodeParser

from llama_index.core.node_parser.relational.hierarchical import (
    HierarchicalNodeParser,
)
from llama_index.core.node_parser.relational.markdown_element import (
    MarkdownElementNodeParser,
)
from llama_index.core.node_parser.relational.unstructured_element import (
    UnstructuredElementNodeParser,
)
from llama_index.core.node_parser.relational.llama_parse_json_element import (
    LlamaParseJsonNodeParser,
)

from llama_index.core import SimpleDirectoryReader
from llama_index.core.node_parser import HierarchicalNodeParser

text="""
深夜的雨声敲打着老式窗棂，林夏在整理祖父遗物时，发现了一封用火漆封缄的信件。信封上褪色的钢笔字迹写着1987年夏至，收件人姓名处被刻意涂抹，只余下林氏宗亲四个模糊的铅字。当她用裁纸刀划开蜡封时，一张泛黄的相片滑落——画面里年轻的祖父站在码头，背景中隐约可见半截被烧焦的船帆，而相片背面用红墨水写着：他们回来了。
阁楼的老式座钟突然敲响十二下，林夏的指尖在相片边缘触到某种凹凸。对着台灯细看，才发现照片的船帆部分竟是用微缩墨水密写的坐标，而信封内衬夹层里还藏着半张船票，目的地是祖父日记里反复提到的幽灵岛。窗外的闪电掠过时，她注意到信纸边缘有被撕去的痕迹，残留的纤维中嵌着几粒与相片背景里相同的黑色砂砾。 林夏将船票与相片上的坐标重叠时，发现它们拼合出三个关键地点：祖父书房暗格里的罗盘、码头废弃仓库的锈蚀铁门，以及镇外被雷劈焦的百年老槐树。当她在罗盘背面刮出第三组数字时，仓库铁门突然在风中发出呻吟般的吱呀声——门缝下塞着张与信封同材质的纸片，上面用祖父的笔迹写着：当砂砾指向月亮，槐树会开口说话。
她冒雨冲向镇外，发现老槐树焦黑的树干上竟嵌着块青铜罗盘，与书房那枚形成镜像。当两枚罗盘指针同时指向月亮时，树皮裂缝中渗出黑色黏液，露出被砂砾封存的黄铜钥匙。钥匙齿纹与相片船帆的密文完全吻合，而仓库铁门后传来祖父录音带播放的沙沙声：...船沉那天，他们带走了真正的...（电流杂音）...砂砾是...（突然中断） 林夏将黄铜钥匙插入仓库铁门的瞬间，生锈的铰链发出刺耳的呻吟。门后弥漫着海盐与铁锈混合的气息，昏黄的灯泡下，祖父的录音机仍在播放断断续续的呓语。她注意到墙角堆放着十几个相同的木箱，每个箱盖都嵌着与相片船帆同材质的黑色砂砾。当钥匙转动到第三圈时，录音机突然发出清晰的咔嗒声——那是祖父年轻时在海军服役的摩斯密码训练录音，但最后三组密码被刻意重复了七遍。
月光透过高窗照在木箱上，砂砾组成的箭头指向中央的保险箱。林夏输入从罗盘和槐树裂缝中拼出的数字时，箱内传来齿轮转动的声响。祖父的日记本静静躺在防潮袋里，最新一页的日期竟是上周，字迹因颤抖而扭曲：...他们每年夏至都会回来取走砂砾...那不是普通的黑曜石...（此处有撕页痕迹）...林家的血脉...
录音带突然自动倒带，播放出二十年前的海浪声与模糊的对话。林夏在保险箱夹层发现半张烧焦的船员名单，与相片背景中船帆的破损处完美吻合。名单末尾的签名处，祖父的名字被红笔圈出，旁边标注着：守门人。当她把所有砂砾集中到月光下时，那些黑色颗粒竟开始缓慢移动，在水泥地上拼出幽灵岛的经纬度坐标——与船票目的地完全一致。 月光下的砂砾突然停止移动，林夏的指尖悬停在那个闪烁的坐标上方。她终于明白，祖父用三十年时间守护的不仅是家族秘密，更是某种正在苏醒的古老契约。当她把所有证据摊开在仓库地板上时，录音机里传来祖父最后的声音：...夏至已至，是时候让真相...（被海浪声淹没）。窗外，雨停了，但林夏知道，这场跨越三十年的谜局才刚刚开始。 （AI生成）
"""

# 加载文档并初始化三级分块解析器
documents = [Document(text=text) for text in text.split("\n")]
print(len(documents))
parser = HierarchicalNodeParser.from_defaults(
    chunk_sizes=[2048, 512, 128]  # 父/中/子节点大小
)
nodes = parser.get_nodes_from_documents(documents)

storage_context = StorageContext.from_defaults()

base_index = VectorStoreIndex(
    nodes,
    storage_context=storage_context,
)

from llama_index.core.retrievers import AutoMergingRetriever

base_retriever = base_index.as_retriever(similarity_top_k=6)
retriever = AutoMergingRetriever(base_retriever, storage_context, verbose=True)

# query_str = "What were some lessons learned from red-teaming?"
# query_str = "Can you tell me about the key concepts for safety finetuning"
query_str = (
    "祖父遗物"
)

nodes = retriever.retrieve(query_str)
base_nodes = base_retriever.retrieve(query_str)
print(len(base_nodes))
print(len(nodes))